import sys
from collections import defaultdict

import backtrader as bt
import pandas as pd


class ProgressObserver(bt.Observer):
    """
    一个极其稳健的观察者，通过定期打印新行来报告进度，以确保最大兼容性。
    """
    lines = ('dummy',)          # 增加一个虚拟line来满足backtrader的内部要求
    plotinfo = dict(plot=False) # 确保这个虚拟line不会被绘制

    def __init__(self):
        self._count = 0
        self._total = 0
        # 将虚拟line指向一个已存在的数据线（例如大盘的收盘价），这不会增加额外计算
        self.lines.dummy = self.data0.close

    def start(self):
        # 使用 len() 来判断数据是否存在，这是最安全的方式，可以避免__bool__类型错误
        if len(self.datas) > 0:
            self._total = len(self.datas[0])
        # 明确告知用户，耗时的初始化已结束，逐日回测开始
        print("\n--- 引擎初始化完成，回测正式开始逐日模拟 ---")

    def next(self):
        # 增加一个安全检查，防止在没有任何数据的情况下（总天数为0）执行后续代码
        if self._total == 0:
            return
            
        self._count += 1
        # 每隔约1年（250个交易日）打印一次进度，避免刷屏
        if self._count % 250 == 0:
            progress = self._count / self._total * 100
            # 获取当前回测的日期
            current_date = self.data0.datetime.date(0)
            print(f"  ...回测已进行到 {current_date}, 进度: {progress:.1f}% ({self._count}/{self._total} 天)")
            sys.stdout.flush() # 强制刷新输出

    def stop(self):
        # 回测结束后，打印最终确认信息
        print(f"--- 每日数据模拟完成，总计 {self._total} 天 ---")

class StockSelectionStrategy(bt.Strategy):
    """
    信号驱动型策略：本策略不再进行任何计算，仅根据外部传入的买入信号执行交易。
    """
    params = (
        ('buy_signals', None), # 接收预计算的买入信号
        ('backtest_start_date', None), # 仍然需要预热期判断
        ('trailing_stop_pct', 0.08), # 跟踪止损的回撤百分比
        ('trade_size', 100), # 新增：每次交易的股数
        ('sell_on_long_upper_shadow_pct', 0.03) # 新增：长上影线卖出阈值
    )

    def __init__(self):
        # 确保信号已传入
        if self.p.buy_signals is None:
            raise ValueError("必须通过参数传入 'buy_signals'")
        
        self.buy_signals = self.p.buy_signals
        self.backtest_start_date = pd.to_datetime(self.p.backtest_start_date)
        
        # --- 引入MA和RSI，用于止损/止盈逻辑 ---
        self.inds = {}
        self.stocks = self.datas[1:] # 获取所有股票数据
        for d in self.stocks:
            self.inds[d] = {
                'ma5': bt.indicators.SimpleMovingAverage(d.close, period=5),
                'ma10': bt.indicators.SimpleMovingAverage(d.close, period=10),
                'ma20': bt.indicators.SimpleMovingAverage(d.close, period=20),
                'rsi': bt.indicators.RSI(d.close, period=14)
            }
        # ---------------------------

        # 创建从股票名称到其数据对象的映射，以便快速查找
        self.stocks_mapping = {d._name: d for d in self.datas[1:]}

        self.trade_log = []
        self.trade_sizes = {}
        # 用于跟踪每笔交易期间的最高价
        self.highest_highs = {}
        # 自己维护一个开放交易的列表，不再依赖 self._trades
        self.open_trades = defaultdict(list)
        self.first_run = True # 标志位，用于首次运行的诊断

        # --- 简化版诊断系统 ---
        self.diag_counters = {
            'total_signals_provided': sum(len(v) for v in self.buy_signals.values()),
            'signals_missed_before_start': 0, # 新增计数器
            'signals_in_period': 0,
            'skipped_already_own': 0,
            'rejected_by_cash': 0,
            'buy_orders_created': 0
        }

    def notify_trade(self, trade):
        if trade.isopen:
            # 当交易开启时，记录其规模和初始最高价（即买入价）
            self.trade_sizes[trade.ref] = trade.size
            self.highest_highs[trade.ref] = trade.price
            # 将新开启的交易加入到我们自己的跟踪列表
            self.open_trades[trade.data].append(trade)

        elif trade.isclosed:
            # 当交易关闭时，清理相关的记录
            original_size = self.trade_sizes.pop(trade.ref, None)
            self.highest_highs.pop(trade.ref, None)
            
            # 从我们自己的跟踪列表中移除已关闭的交易
            closed_ref = trade.ref
            trade_list = self.open_trades.get(trade.data, [])
            self.open_trades[trade.data] = [t for t in trade_list if t.ref != closed_ref]
            # 如果列表为空，则从字典中删除该键
            if not self.open_trades[trade.data]:
                self.open_trades.pop(trade.data, None)

            # 使用 .get 避免在并发场景下的KeyError - 修正: 此处应使用pop出的值
            sell_price = (trade.pnl / original_size) + trade.price if original_size and original_size != 0 else 0
            buy_cost = trade.price * original_size if original_size else 0
            pnl_ratio = trade.pnl / buy_cost if buy_cost != 0 else 0.0
            # 对于coc=True设置，交易实际是在当天收盘时执行的，但Backtrader记录的是下一个bar的时间
            # 因此我们需要将日期向前调整一天，以反映实际的交易执行日期
            buy_dt = bt.num2date(trade.dtopen)
            sell_dt = bt.num2date(trade.dtclose)
            
            # 修正买入和卖出日期（向前调整一天）
            actual_buy_dt = buy_dt - pd.Timedelta(days=1)
            actual_sell_dt = sell_dt - pd.Timedelta(days=1)
            
            log_entry = {
                'symbol': trade.data._name,
                'buy_datetime': actual_buy_dt.strftime('%Y-%m-%d %H:%M:%S'),
                'buy_price': trade.price,
                'sell_datetime': actual_sell_dt.strftime('%Y-%m-%d %H:%M:%S'),
                'sell_price': sell_price,
                'holding_days': trade.barlen,
                'pnl': trade.pnl,
                'pnl_ratio': f"{pnl_ratio * 100:.2f}%"}
            self.trade_log.append(log_entry)

    def stop(self):
        # --- 新的诊断报告 ---
        print("\n--- 策略诊断回溯 (信号驱动模型) ---")
        total_signals = self.diag_counters['total_signals_provided']
        missed_signals = self.diag_counters['signals_missed_before_start']
        in_period = self.diag_counters['signals_in_period']
        own_skip = self.diag_counters['skipped_already_own']
        cash_rej = self.diag_counters['rejected_by_cash']
        bought = self.diag_counters['buy_orders_created']

        print(f"1. Pandas海选阶段共生成买入信号: {total_signals}")
        print(f"   - (过滤) 因数据未对齐/新股等原因，在回测循环开始前被错过: {missed_signals}")
        print(f"---------------------------------------------------")
        print(f"2. 进入实际交易执行环节的信号: {in_period}")
        print(f"   - (跳过) 因已持有该股: {own_skip}")
        print(f"   - (拒绝) 因资金不足: {cash_rej}")
        print(f"   - (成交) 成功生成买入订单: {bought}")
        print(f"===================================================")

        # 验证计数器逻辑
        accounted_for = own_skip + cash_rej + bought
        if in_period != accounted_for:
             print(f"   -> 警告: 漏斗下游计数错误! 有效信号({in_period}) != 各项结果之和({accounted_for})")
        
        # 增加总体验证
        if total_signals != missed_signals + in_period:
            print(f"   -> 警告: 总体验证失败! 总信号({total_signals}) != 错过信号({missed_signals}) + 有效信号({in_period})")

        print("----------------------\n")

    def next(self):
        current_date = self.datetime.date(0)

        # 在第一次进入next方法时，进行一次性诊断
        if self.first_run:
            self.first_run = False
            missed_count = 0
            for signal_date, signals in self.buy_signals.items():
                if signal_date < current_date:
                    missed_count += len(signals)
            self.diag_counters['signals_missed_before_start'] = missed_count

        # --- 1. 遍历所有持仓，检查卖出条件 ---
        # 创建一个需要关闭的交易列表，以避免在迭代时修改持仓
        trades_to_close = []
        # 遍历自己维护的 open_trades 字典
        for data, trade_list in list(self.open_trades.items()):
            for t in trade_list:
                # --- 止盈逻辑 ---
                # 步骤1.1: 更新期间最高价
                current_high = self.highest_highs.get(t.ref, t.price)
                self.highest_highs[t.ref] = max(current_high, data.high[0])
                # 步骤1.2: 计算跟踪止盈触发价
                trailing_stop_price = self.highest_highs[t.ref] * (1.0 - self.p.trailing_stop_pct)
                
                sell_cond_1 = data.close[0] < trailing_stop_price

                # --- 止损逻辑 ---
                is_losing = data.close[0] < t.price
                below_ma10 = data.close[0] < self.inds[data]['ma10'][0]
                
                sell_cond_2 = is_losing and below_ma10

                # --- 新增：固定比例止损 ---
                loss_pct = (data.close[0] / t.price) - 1.0
                sell_cond_3 = loss_pct < -0.05

                # --- 新增：盈利保护止损（盈利状态下，跌破20日线卖出） ---
                # is_profitable = data.close[0] > t.price
                # below_ma20 = data.close[0] < self.inds[data]['ma20'][0]
                # sell_cond_4 = is_profitable and below_ma20

                # --- 新增：基于长上影线的卖出信号 ---
                long_upper_shadow = (data.high[0] - data.close[0]) / data.high[0] > self.p.sell_on_long_upper_shadow_pct if data.high[0] > 0 else False
                sell_cond_5 = long_upper_shadow

                # --- 最终判断 ---
                if sell_cond_1 or sell_cond_2 or sell_cond_3 or sell_cond_5:
                    trades_to_close.append(data)
        
        # 统一执行平仓操作
        for d in set(trades_to_close): # 使用set去重，防止重复平仓
            self.close(data=d)

        # --- 2. 预热期"门卫" ---
        if current_date < self.backtest_start_date.date():
            return

        # --- 3. 根据预计算信号执行买入 ---
        signals_today = self.buy_signals.get(current_date)
        
        if not signals_today:
            return # 今天没有信号

        self.diag_counters['signals_in_period'] += len(signals_today)

        for symbol in signals_today:
            d = self.stocks_mapping.get(symbol)
            if d is None: # 容错：如果信号中的股票没有被加载
                continue
            
            # 检查是否已持有
            if self.getposition(d):
                self.diag_counters['skipped_already_own'] += 1
                continue
            
            # 检查资金是否充足
            size = self.p.trade_size # 使用配置的交易股数
            required_cash = size * d.close[0]
            if self.broker.get_cash() < required_cash:
                self.diag_counters['rejected_by_cash'] += 1
                continue

            self.buy(data=d, size=size)
            self.diag_counters['buy_orders_created'] += 1 